#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include "mpi.h"
#include "lib.h"

/* Aloca memoria */
void allocData(PPMImage *img, char* type, int width, int height, int maxVal) {
    strcpy(img->type, type);
    img->width = width;
    img->height = height;
    img->maxVal = maxVal;

    img->r = (uchar *)malloc(width * height * sizeof(uchar));
    img->g = (uchar *)malloc(width * height * sizeof(uchar));
    img->b = (uchar *)malloc(width * height * sizeof(uchar));
    img->k = (int *)malloc(width * height * sizeof(int));
}

/* Libera memoria */
void freeData(PPMImage *img) {
    free(img->r);
    free(img->g);
    free(img->b);
    free(img->k);
}

/* Cria estrutura PPMImage */
void readImage(PPMImage *imgIn, PPMImage *imgOut, PPMImage *vecIn, PPMImage *vecOut, char *in){
    FILE *input;
    int i, j, k, z;
    char type[2];
    int maxVal, width, height, blocWidth, blocHeight;

    input = fopen(in, "r");

    if (input == NULL)
        printf("error\n");
    else {
        /* Le header */
        // Le tipo
        fscanf(input, "%s", type);
        // Pula fim da linha
        while (getc(input) != '\n');

        // Pula comentario
        while (getc(input) == '#') {
            while (getc(input) != '\n');
        }

        // Volta um caracter
        fseek(input, -1, SEEK_CUR);

        // Le tamanho
        fscanf(input, "%d", &width);
        fscanf(input, "%d", &height);
        fscanf(input, "%d", &maxVal);

        allocData(imgIn, type, width, height, maxVal);
        allocData(imgOut, type, width, height, maxVal);

        // Le dados e separa em canais
        for(i = 0; i < width * height ; i++) {
            fscanf(input, "%hhu %hhu %hhu", &(imgIn->r[i]), &(imgIn->g[i]), &(imgIn->b[i]));
        }

        // Recalcula dimensoes das imagens menores
        blocWidth = width / (hosts / 2);
        blocHeight = height / (hosts / 2);

        // Aloca memoria para as imagens menores
        for(i = 0; i < hosts ; i++) {
            allocData(&vecIn[i], type, blocWidth + RADIUS, blocHeight + RADIUS, maxVal);
            allocData(&vecOut[i], type, blocWidth + RADIUS, blocHeight + RADIUS, maxVal);
        }

        // Cria as imagens menores
        z = 0;
        for(i = 0; i < blocHeight + RADIUS; i++) {
            for(j = 0; j < blocWidth + RADIUS; j++) {
                k = (i * width) + j;
                vecIn[0].r[z] = imgIn->r[k];
                vecIn[0].g[z] = imgIn->g[k];
                vecIn[0].b[z] = imgIn->b[k];
                vecIn[0].k[z] = k;
                
                k = (blocWidth - RADIUS) + ((i * width) + j);
                vecIn[1].r[z] = imgIn->r[k];
                vecIn[1].g[z] = imgIn->g[k];
                vecIn[1].b[z] = imgIn->b[k];
                vecIn[1].k[z] = k;
                
                k = ((blocHeight - RADIUS) * width) + ((i * width) + j);
                vecIn[2].r[z] = imgIn->r[k];
                vecIn[2].g[z] = imgIn->g[k];
                vecIn[2].b[z] = imgIn->b[k];
                vecIn[2].k[z] = k;
                
                k = ((blocHeight - RADIUS) * width) + (blocWidth - RADIUS) + ((i * width) + j);
                vecIn[3].r[z] = imgIn->r[k];
                vecIn[3].g[z] = imgIn->g[k];
                vecIn[3].b[z] = imgIn->b[k];
                vecIn[3].k[z] = k;

                z++;
            }
        }
    }
    fclose(input);
}

/* Cria arquivo com o resultado */
void saveImage(PPMImage *img, PPMImage *vec, char *out) {
    FILE *output;
    int i, j;

    // Escreve header
    output = fopen(out, "w");
    fprintf(output, "%s\n", img->type);
    fprintf(output, "%d %d\n", img->width, img->height);
    fprintf(output, "%d\n", img->maxVal);

    // Refaz a imagem grande com as menores
    for (i = 0; i < vec->height; i++) {
        for (j = 0; j < vec->width; j++) {
            img->r[vec[0].k[(i * (vec->width)) + j]] = vec[0].r[(i * (vec->width)) + j];
            img->g[vec[0].k[(i * (vec->width)) + j]] = vec[0].g[(i * (vec->width)) + j];
            img->b[vec[0].k[(i * (vec->width)) + j]] = vec[0].b[(i * (vec->width)) + j];

            img->r[vec[1].k[(i * (vec->width)) + j + RADIUS]] = vec[1].r[(i * (vec->width)) + j + RADIUS];
            img->g[vec[1].k[(i * (vec->width)) + j + RADIUS]] = vec[1].g[(i * (vec->width)) + j + RADIUS];
            img->b[vec[1].k[(i * (vec->width)) + j + RADIUS]] = vec[1].b[(i * (vec->width)) + j + RADIUS];

            img->r[vec[2].k[((i + RADIUS) * (vec->width)) + j]] = vec[2].r[((i + RADIUS) * (vec->width)) + j];
            img->g[vec[2].k[((i + RADIUS) * (vec->width)) + j]] = vec[2].g[((i + RADIUS) * (vec->width)) + j];
            img->b[vec[2].k[((i + RADIUS) * (vec->width)) + j]] = vec[2].b[((i + RADIUS) * (vec->width)) + j];

            img->r[vec[3].k[((i + RADIUS) * (vec->width)) + j + RADIUS]] = vec[3].r[((i + RADIUS) * (vec->width)) + j + RADIUS];
            img->g[vec[3].k[((i + RADIUS) * (vec->width)) + j + RADIUS]] = vec[3].g[((i + RADIUS) * (vec->width)) + j + RADIUS];
            img->b[vec[3].k[((i + RADIUS) * (vec->width)) + j + RADIUS]] = vec[3].b[((i + RADIUS) * (vec->width)) + j + RADIUS];
        }
    }

    // Escreve dados
    for (i = 0; i < img->height * img->width; i++) {
        fprintf(output, "%hhu %hhu %hhu ", img->r[i], img->g[i], img->b[i]);
    }

    fclose(output);
}

int main(int argc, char const *argv[]) {
    MPI_Status status;
    MPI_Init(&argc, &argv);
    int processId;

    MPI_Comm_rank(MPI_COMM_WORLD, &processId);


    if (processId == MASTER) {
        PPMImage imgIn, imgOut, vecIn[HOSTS], vecOut[HOSTS];
        char in[20], out[20];

        strcpy(in, argv[1]);
        strcpy(out, argv[2]);

        readImage(&imgIn, &imgOut, vecIn, vecOut, in);

        for (i = 0; i < HOSTS; i++) {
            // Envia pros nos
            MPI_Send(&(vecIn[i].height), 1, MPI_INT, i + 1, 0, MPI_COMM_WORLD);
            MPI_Send(&(vecIn[i].width), 1, MPI_INT, i + 1, 1, MPI_COMM_WORLD);
            MPI_Send(&(vecIn[i].maxVal), 1, MPI_INT, i + 1, 2, MPI_COMM_WORLD);
            MPI_Send(&(vecIn[i].r[0]), vecIn.width * vecIn.height, MPI_UNSIGNED_CHAR, i + 1, 3, MPI_COMM_WORLD);
            MPI_Send(&(vecIn[i].g[0]), vecIn.width * vecIn.height, MPI_UNSIGNED_CHAR, i + 1, 4, MPI_COMM_WORLD);
            MPI_Send(&(vecIn[i].b[0]), vecIn.width * vecIn.height, MPI_UNSIGNED_CHAR, i + 1, 5, MPI_COMM_WORLD);
        }

        for (i = 0; i < HOSTS; i++) {
            // Recebe
            MPI_Recv(&(vecOut.r[0]), vecOut.width * vecOut.height, MPI_UNSIGNED_CHAR, i + 1, 3, MPI_COMM_WORLD, &status);
            MPI_Recv(&(vecOut.g[0]), vecOut.width * vecOut.height, MPI_UNSIGNED_CHAR, i + 1, 4, MPI_COMM_WORLD, &status);
            MPI_Recv(&(vecOut.b[0]), vecOut.width * vecOut.height, MPI_UNSIGNED_CHAR, i + 1, 5, MPI_COMM_WORLD, &status);
        }
        saveImage(&imgOut, vecOut, out);

        for (i = 0; i < HOSTS; i++) {
            freeData(&vecIn[i]);
            freeData(&vecOut[i]);  
        }

        freeData(&imgIn);
        freeData(&imgOut); 
    } else if(processId > 0) {
        PPMImage vecIn, vecOut;
        int height, width, maxVal;

        MPI_Recv(&(width), 1, MPI_INT, MASTER, 0, MPI_COMM_WORLD, &status);
        MPI_Recv(&(height), 1, MPI_INT, MASTER, 1, MPI_COMM_WORLD, &status);
        MPI_Recv(&(maxVal), 1, MPI_INT, MASTER, 2, MPI_COMM_WORLD, &status);

        allocData(&vecIn, 'P3', width, height, maxVal);
        allocData(&vecOut, 'P3', width, height, maxVal);
        
        MPI_Recv(&(vecIn.r[0]), width * height, MPI_UNSIGNED_CHAR, MASTER, 3, MPI_COMM_WORLD, &status);
        MPI_Recv(&(vecIn.g[0]), width * height, MPI_UNSIGNED_CHAR, MASTER, 4, MPI_COMM_WORLD, &status);
        MPI_Recv(&(vecIn.b[0]), width * height, MPI_UNSIGNED_CHAR, MASTER, 5, MPI_COMM_WORLD, &status);

        launchCuda(&vecIn,&vecOut);

        MPI_Send(&(vecOut[i].r[0]), vecOut.width * vecOut.height, MPI_UNSIGNED_CHAR, MASTER, 3, MPI_COMM_WORLD);
        MPI_Send(&(vecOut[i].g[0]), vecOut.width * vecOut.height, MPI_UNSIGNED_CHAR, MASTER, 4, MPI_COMM_WORLD);
        MPI_Send(&(vecOut[i].b[0]), vecOut.width * vecOut.height, MPI_UNSIGNED_CHAR, MASTER, 5, MPI_COMM_WORLD);
    }
    
    return 0;
}

